#include <gtest/gtest.h>

#include <future>

#include "scheduler.hpp"

TEST(OldSchedulerTest, Init) {
  OldScheduler scheduler(3, 1);
  EXPECT_EQ(scheduler.cpu_worker_num(), 3);
  EXPECT_EQ(scheduler.gpu_worker_num(), 1);
}

TEST(OldSchedulerTest, SubmitSingleTask) {
  OldScheduler scheduler;

  auto future = scheduler.submit_task(
      0, [](int num_a, int numb_b) { return num_a + numb_b; }, 3, 4);

  EXPECT_EQ(future.get(), 7);
}

TEST(OldSchedulerTest, SubmitMultiTask) {
  const int LOOP = 1000;

  OldScheduler scheduler;
  std::vector<std::future<int>> rets;
  rets.reserve(LOOP);

  for (int i = 0; i < LOOP; ++i) {  // NOLINT
    rets.emplace_back(scheduler.submit_task(
        i % scheduler.cpu_worker_num(),
        [](int num_a, int numb_b) { return num_a + numb_b; }, i, i));
  }

  for (int i = 0; i < LOOP; ++i) {  // NOLINT
    EXPECT_EQ(rets[i].get(), i + i);
  }
}
